{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from osgeo import gdal\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from keras.layers import *\n",
    "from keras.models import Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_tiff(fn):\n",
    "    \"\"\"\n",
    "        inputs: tiff filename\n",
    "        pouputs: image data\n",
    "    \"\"\"\n",
    "    dataset = gdal.Open(fn,gdal.GA_ReadOnly)\n",
    "    im_width = dataset.RasterXSize\n",
    "    im_height = dataset.RasterYSize\n",
    "    im_bands = dataset.RasterCount\n",
    "    im_data = dataset.ReadAsArray(0,0,im_width,im_height)\n",
    "    if len(im_data.shape)==2:\n",
    "        return im_data\n",
    "    else:\n",
    "        return im_data.transpose([1,2,0])\n",
    "def features(dt):\n",
    "    \"\"\"\n",
    "        inputs: images\n",
    "        pouputs: images with feature bands\n",
    "    \"\"\"\n",
    "    Blue = dt[:,:,1].reshape(dt.shape[:2]+(1,))\n",
    "    Green = dt[:,:,2].reshape(dt.shape[:2]+(1,))\n",
    "    Red = dt[:,:,3].reshape(dt.shape[:2]+(1,))\n",
    "    NIR = dt[:,:,4].reshape(dt.shape[:2]+(1,))\n",
    "    SWIR1 = dt[:,:,5].reshape(dt.shape[:2]+(1,))\n",
    "    mndwi = dt[:,:,5].reshape(dt.shape[:2]+(1,))\n",
    "    leb = dt[:,:,6].reshape(dt.shape[:2]+(1,))\n",
    "    dem = dt[:,:,7].reshape(dt.shape[:2]+(1,))\n",
    "    return np.concatenate([Blue,Green,Red,NIR,ndvi,mndwi,leb,dem],axis=-1)\n",
    "def get_input(label, im):\n",
    "    \"\"\"\n",
    "        inputs: None\n",
    "        outputs: model\n",
    "    \"\"\"\n",
    "    features_im = features(im)\n",
    "    y = label[3:-3,3:-3]\n",
    "    water_pixel_num = (y==1).sum()\n",
    "    other_pixel_num = (y==0).sum()\n",
    "    if water_pixel_num<other_pixel_num:\n",
    "        loc_x = np.where(y!=1)[0]+3\n",
    "        loc_y = np.where(y!=1)[1]+3\n",
    "        loc_length = len(loc_x)\n",
    "        loc_choice = np.random.choice(np.arange(loc_length),min(int(water_pixel_num*5),loc_length),replace=False)\n",
    "        loc_x = loc_x[loc_choice]#np.random.choice(np.where(t!=1)[0],water_pixel_num,replace=False)\n",
    "        loc_y = loc_y[loc_choice]#np.random.choice(np.where(t!=1)[1],water_pixel_num,replace=False)\n",
    "        label[loc_x,loc_y] = 2\n",
    "    elif water_pixel_num>other_pixel_num:\n",
    "        label[np.where(y==0)[0]+3,np.where(y==0)[1]+3]=2\n",
    "        loc_x = np.where(y!=0)[0]+3\n",
    "        loc_y = np.where(y!=0)[1]+3\n",
    "        label[loc_x,loc_y] = 0\n",
    "        loc_length = len(loc_x)\n",
    "        loc_choice = np.random.choice(np.arange(loc_length),other_pixel_num,replace=False)\n",
    "        loc_x = loc_x[loc_choice]#np.random.choice(np.where(t!=1)[0],water_pixel_num,replace=False)\n",
    "        loc_y = loc_y[loc_choice]#np.random.choice(np.where(t!=1)[1],water_pixel_num,replace=False)\n",
    "        label[loc_x,loc_y] = 1\n",
    "    else:\n",
    "        label = 2-label\n",
    "    loc_x = np.where(label[3:-3,3:-3]!=0)[0]\n",
    "    loc_y = np.where(label[3:-3,3:-3]!=0)[1]\n",
    "    pixel_num = len(loc_x)\n",
    "    #print(pixel_num)\n",
    "    im_train = np.zeros((pixel_num,7,7,8))\n",
    "    lb_train = np.zeros((pixel_num,1,1,2))\n",
    "    num = 0\n",
    "    for i,j in zip(loc_x,loc_y):\n",
    "        try:\n",
    "            im_train[num] = features_im[i:i+7,j:j+7,:]\n",
    "        except:\n",
    "            print(i,j)\n",
    "        if label[i+3,j+3] == 1:\n",
    "            lb_train[num,:,:,0] = 1\n",
    "        elif label[i+3,j+3] == 2:\n",
    "            lb_train[num,:,:,1] = 1\n",
    "        num+=1\n",
    "    return im_train,lb_train\n",
    "def CNN():\n",
    "    \"\"\"\n",
    "        inputs: None\n",
    "        outputs: model\n",
    "    \"\"\"\n",
    "    inputs = Input((7,7,8))\n",
    "    x = Conv2D(16, (3, 3), activation = 'relu')(inputs)\n",
    "    x = Conv2D(32, (3, 3), activation = 'relu')(x)\n",
    "    x = concatenate([x,inputs[:,2:-2,2:-2,:]])\n",
    "    x = Conv2D(64, (3, 3), activation = 'relu')(x)\n",
    "    x = concatenate([x,inputs[:,3:-3,3:-3,:]])\n",
    "    x = Conv2D(128, (1, 1), activation = 'relu')(x)\n",
    "    outputs = Conv2D(2,(1,1),activation='sigmoid')(x)\n",
    "    model = Model(inputs=inputs, outputs=outputs)\n",
    "    return model\n",
    "def generate_inputs(label_fns, im_fns):\n",
    "    \"\"\"\n",
    "        inputs: filename of labels and images\n",
    "        outputs: None\n",
    "    \"\"\"\n",
    "    for i, j in zip(label_fns, im_fns):\n",
    "        im = read_tiff(j)\n",
    "        label = read_tiff(i)\n",
    "        label = label/label.max()\n",
    "        yield get_input(label, im)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = CNN()\n",
    "fns = os.listdir(r'[folder]')\n",
    "######################################\n",
    "# folder of labels\n",
    "######################################\n",
    "label_fns = [os.path.join(r'[folder]',i) for i in fns]\n",
    "######################################\n",
    "# folder of labels\n",
    "######################################\n",
    "im_fns = [os.path.join(r'[folder]','im'+i[4:]) for i in fns]\n",
    "######################################\n",
    "# folder of images\n",
    "######################################\n",
    "epochs = {}\n",
    "for ep in range(10):\n",
    "    print('epoch:',ep+1)\n",
    "    epochs[str(ep+1)] = {'loss':[],'acc':[]}\n",
    "    with open(r'[filename]','a+') as f:\n",
    "        ######################################\n",
    "        # log filename\n",
    "        ######################################\n",
    "        f.write('epoch:'+str(ep+1)+'\\n')\n",
    "    num=1\n",
    "    for i, j in generate_inputs(label_fns, im_fns):\n",
    "        print('file',num)\n",
    "        if i.shape[0]:\n",
    "            history=model.fit(i,j,batch_size=4096,epochs = 1,shuffle = True)\n",
    "            with open(r'[filename]','a+') as f:\n",
    "                ######################################\n",
    "                # log filename\n",
    "                ######################################\n",
    "                f.write(str(history.history)+'\\n')\n",
    "            epochs[str(ep+1)]['loss'].append(history.history['loss'][0])\n",
    "            epochs[str(ep+1)]['acc'].append(history.history['acc'][0])\n",
    "        num+=1\n",
    "    \n",
    "    model.save(r'[folder]'+'\\\\model_'+str(ep)+'.h5')\n",
    "    ######################################\n",
    "    # folder to save models\n",
    "    ######################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
